# 1.导入streamlit库
import streamlit as st
import ollama
# PS：ollama默认访问的是本地ollama服务，所以需要先启动ollama服务：http://127.0.0.1:11434
import os
os.environ["OLLAMA_HOST"]="http://127.0.0.1:11434"

# 2.设置页面标题并添加分割线
st.title("聊天机器人")
# 添加分割线
st.divider()

# 3.先判断memory是否在session_state中，如果不在就创建对象
if "message" not in st.session_state:
    # 先添加一条ai提示的记忆，后续都是实时追加最新消息
    st.session_state["message"]=[{
        "role":"assistant",
        "content":"你好，我是聊天机器人，有什么我可以帮助你的吗？"
    }]
# 4.聊天窗口
# （1）AI先说话
for message in st.session_state["message"]:
    st.chat_message(message["role"]).write(message["content"])
# (2)用户问问题
prompt=st.chat_input()


def get_chat_response(messages):
    print(os.getenv("OLLAMA_HOST"))
    # 发送请求给本地大模型，获取结果
    result=ollama.chat(
        model='qwen3:32b',  # 修正模型名称
        messages=messages,
    )
    # 返回结果
    return result["message"]["content"]


if prompt:
    # 如果用户输入apikey，那就展示用户提问的问题到页面并添加记忆
    st.chat_message("user").write(prompt)
    st.session_state["message"].append({
        "role":"user",
        "content":prompt
    })
    # (3)AI回答
    with st.spinner("正在思考中..."):
        res=get_chat_response(st.session_state.message[-30:])
    # 如果获取到了ai的回答，就展示到页面，并添加记忆
    st.chat_message("assistant").write(res)
    st.session_state["message"].append({
        "role":"assistant",
        "content":res
    })
